package cn.st4rlight.util.excel;

import static cn.st4rlight.util.excel.DynamicExcelUtil.MAX_ROWS_PER_SHEET;

import java.io.ByteArrayOutputStream;
import java.util.List;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.ListUtils;
import org.apache.commons.lang3.time.StopWatch;

import com.alibaba.excel.EasyExcelFactory;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.builder.ExcelWriterSheetBuilder;
import com.alibaba.excel.write.handler.WriteHandler;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.google.common.collect.Lists;

import cn.st4rlight.util.excel.context.SingleSheetContext;
import cn.st4rlight.util.excel.strategy.AutoWidthStrategy;
import cn.st4rlight.util.exception.ServiceException;
import cn.st4rlight.util.value.SafeUtil;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

/**
 * @author st4rlight <st4rlight@163.com>
 * Created on 2024-04-15
 */
@Slf4j
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class WriteExcelUtil {

    /**
     * 构建只有单个sheet的excel文件
     *
     * @param context 上下文参数
     * @return 所生成的excel文件字节数组
     */
    public static <T> byte[] writeSingleSheetExcel(SingleSheetContext<T> context) {
        return WriteExcelUtil.writeSingleSheetExcel(context, Lists.newArrayList());
    }

    /**
     * 构建只有单个sheet的excel文件
     *
     * @param context 上下文参数
     * @param handlerList 定制化写出策略列表
     * @return 所生成的excel文件字节数组
     */
    public static <T> byte[] writeSingleSheetExcel(SingleSheetContext<T> context, List<WriteHandler> handlerList) {
        return WriteExcelUtil.writeMultiSheetExcel(Lists.newArrayList(context), handlerList);
    }


    /**
     * 构建有多个sheet的excel文件，顺序即按contextList中的顺序
     *
     * @param contextList 上下文参数信息列表
     * @return 所生成的excel文件字节数组
     */
    public static <T> byte[] writeMultiSheetExcel(List<SingleSheetContext<T>> contextList) {
        return WriteExcelUtil.writeMultiSheetExcel(contextList, Lists.newArrayList());
    }

    /**
     * 构建有多个sheet的excel文件，顺序即按contextList中的顺序
     *
     * @param contextList 上下文参数信息列表
     * @param handlerList 定制化写出策略列表
     * @return 所生成的excel文件字节数组
     */
    public static <T> byte[] writeMultiSheetExcel(List<SingleSheetContext<T>> contextList, List<WriteHandler> handlerList) {
        if (CollectionUtils.isEmpty(contextList)) {
            throw ServiceException.of("用于构建excel的上下参数列表不能为空");
        }

        // 拆分超过100w的context
        List<SingleSheetContext<T>> newContextList = WriteExcelUtil.splitLargeContext(contextList);
        log.info("writeMultiSheetExcel sheet数量: {}", newContextList.size());

        // 开始处理excel导出
        StopWatch stopWatch = StopWatch.createStarted();
        try(
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            ExcelWriter excelWriter = EasyExcelFactory.write(bos).build()
        ) {
            // 逐个转成sheet
            for (int i = 0; i < newContextList.size(); i++) {
                log.info("writeMultiSheetExcel {} / {}", i + 1, newContextList.size());
                SingleSheetContext<T> param = newContextList.get(i);
                WriteSheet writeSheet = writeSingleSheet(i, param, handlerList);
                // 写出当前sheet到输出流
                excelWriter.write(param.getDtoList(), writeSheet);
            }

            // 内部会关闭流
            excelWriter.finish();
            log.info("writeMultiSheetExcel完成 耗时: {}ms, size: {}", stopWatch.getTime(), bos.size());
            return bos.toByteArray();
        } catch (Exception ex) {
            log.error("动态导出excel文件异常", ex);
            throw ServiceException.of(ex.getMessage());
        }
    }

    /**
     * 写出参数到单个sheet
     *
     * @param sheetIndex sheet在excel文件中的下标
     * @param context 上下文参数
     * @param handlerList 定制化写出策略列表
     * @return 单个sheet
     */
    private static <T> WriteSheet writeSingleSheet(
        int sheetIndex, SingleSheetContext<T> context, List<WriteHandler> handlerList
    ) {
        ExcelWriterSheetBuilder builder = EasyExcelFactory.writerSheet(sheetIndex, context.getSheetName())
                .excludeColumnFieldNames(SafeUtil.of(context.getExcludeFieldNames()))
                .registerWriteHandler(new AutoWidthStrategy())
                .head(context.getHeadClz());

        // 定制化导出策略
        if (CollectionUtils.isNotEmpty(handlerList)) {
            handlerList.forEach(builder::registerWriteHandler);
        }

        return builder.build();
    }

    /**
     * 单个sheet最多100w行，拆分context
     */
    public static <T> List<SingleSheetContext<T>> splitLargeContext(List<SingleSheetContext<T>> contextList) {
        List<SingleSheetContext<T>> newContextList = Lists.newArrayList();

        for (SingleSheetContext<T> context : contextList) {
            List<T> dtoList = context.getDtoList();
            if (dtoList.size() <= MAX_ROWS_PER_SHEET) {
                newContextList.add(context);
                continue;
            }

            // 超过100w行拆分成多个sheet
            List<List<T>> partitionList = Lists.partition(dtoList, MAX_ROWS_PER_SHEET);
            for (List<T> partition : partitionList) {
                SingleSheetContext<T> newContext = new SingleSheetContext<T>()
                        .setDtoList(partition)
                        .setHeadClz(context.getHeadClz())
                        .setSheetName(context.getSheetName());
                newContextList.add(newContext);
            }
        }

        log.info("splitLargeContext origin size: {}, result size: {}", contextList.size(), newContextList.size());
        return newContextList;
    }

    /**
     * 简单构建一个excel文件，没有指定header时默认使用字段名
     * 注意：类字段必须有getter方法才行
     *
     * @return 所生成的excel文件字节数组
     */
    public static <T> byte[] easyWriteExcel(List<T> dataList, Class<T> dataClz) {
        List<List<T>> partitionList = ListUtils.partition(dataList, MAX_ROWS_PER_SHEET);

        // 开始处理excel导出
        StopWatch stopWatch = StopWatch.createStarted();
        try (
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            ExcelWriter excelWriter = EasyExcelFactory.write(bos).build()
        ) {
            // 遍历参数列表，逐个转成sheet
            for (int i = 0; i < partitionList.size(); i++) {
                log.info("easyWriteExcel write: {} / {}", i + 1, partitionList.size());
                String sheetName = "sheet" + (i + 1);
                SingleSheetContext<T> param = new SingleSheetContext<T>()
                        .setSheetName(sheetName)
                        .setHeadClz(dataClz);
                WriteSheet writeSheet = writeSingleSheet(i, param, Lists.newArrayList());
                // 写入当前sheet到输出流
                excelWriter.write(partitionList.get(i), writeSheet);
            }

            // 内部会关闭流
            excelWriter.finish();
            log.info("easyWriteExcel完成 耗时: {}ms, size: {}", stopWatch.getTime(), bos.size());
            return bos.toByteArray();
        } catch (Exception ex) {
            log.info("", ex);
            throw ServiceException.of(ex);
        }
    }
}
